import pandas as pd
import numpy as np
import statsmodels.formula.api as smf
import matplotlib.pyplot as plt
import seaborn as sns
import os
from datetime import datetime, timedelta

def generate_financial_data():
    """生成包含三个数据表的Excel文件:日期表、现金流表和指数表"""
    # 获取当前日期作为基准日
    date_today = datetime.now()
    
    # 生成 date_1 sheet 数据（最近30天日期及星期）
    def create_date1_sheet():
        dates = [date_today - timedelta(days=i) for i in range(30)]
        weekdays = [d.strftime('%A') for d in dates]
        return pd.DataFrame({'Date': dates, 'Weekday': weekdays})
    
    # 生成 moneyflows sheet 数据（最近100天随机现金流）
    def create_moneyflows_sheet():
        dates = [date_today - timedelta(days=i) for i in range(100)]
        cash_flows = np.random.randint(-1000, 1000, size=100)
        return pd.DataFrame({'Date': dates, 'Cash_Flow': cash_flows})
    
    # 生成 index_daily sheet 数据（最近365天随机指数值）
    def create_index_sheet():
        dates = [date_today - timedelta(days=i) for i in range(365)]
        index_values = np.random.randint(1000, 10000, size=365)
        return pd.DataFrame({'Date': dates, 'Index_Value': index_values})
    
    # 创建Excel写入对象
    with pd.ExcelWriter('financial_data.xlsx') as writer:
        # 将各sheet写入文件
        create_date1_sheet().to_excel(writer, sheet_name='date_1', index=False)
        create_moneyflows_sheet().to_excel(writer, sheet_name='moneyflows', index=False)
        create_index_sheet().to_excel(writer, sheet_name='index_daily', index=False)
    
    print("financial_data.xlsx 文件已生成,包含以下三个sheet:")
    print("- date_1:最近30天的日期及对应星期")
    print("- moneyflows:最近100天的随机现金流数据")
    print("- index_daily:最近365天的随机指数值数据")

def analyze_financial_data():
    """读取Excel文件并进行数据分析、回归建模和可视化"""
    # 定义文件路径
    file_path = './financial_data.xlsx'
    
    # 检查文件是否存在
    if not os.path.exists(file_path):
        print(f"错误：文件 {file_path} 不存在。请先运行generate_financial_data()生成文件。")
        return
    
    # 1. 使用pandas读取xlsx文件
    try:
        df_date1 = pd.read_excel(file_path, sheet_name='date_1')
        df_moneyflows = pd.read_excel(file_path, sheet_name='moneyflows')
        df_index_daily = pd.read_excel(file_path, sheet_name='index_daily')
    except Exception as e:
        print(f"读取文件时发生错误：{e}")
        return
    
    # 2. 合并数据
    # 先合并date_1和moneyflows（假设ts_code和trade_date为共同列）
    # 注意：原始生成代码未生成ts_code列，这里添加模拟数据
    df_date1['ts_code'] = '000001.SZ'  # 添加模拟股票代码
    df_moneyflows['ts_code'] = '000001.SZ'
    df_index_daily['ts_code'] = 'INDEX'  # 指数使用不同代码
    
    # 合并date_1和moneyflows
    df_merged = pd.merge(df_date1, df_moneyflows, on=['ts_code', 'Date'], how='inner')
    df_merged = df_merged.rename(columns={'Date': 'trade_date'})  # 重命名日期列为trade_date
    
    # 再合并index_daily（左连接）
    df_index_daily = df_index_daily.rename(columns={'Date': 'trade_date'})
    df_merged = pd.merge(df_merged, df_index_daily, on=['trade_date'], how='left', suffixes=('', '_index'))
    
    # 3. 数据处理和分析
    # 计算涨跌幅和指数涨跌幅（使用Cash_Flow代替closes进行演示）
    df_merged['zd_close'] = round((df_merged['Cash_Flow'] - df_merged['Cash_Flow'].shift(1)) / df_merged['Cash_Flow'].shift(1), 2)
    df_merged['hz_close'] = round((df_merged['Index_Value'] - df_merged['Index_Value'].shift(1)) / df_merged['Index_Value'].shift(1), 2)
    
    # 处理缺失值
    df_merged = df_merged.dropna(subset=['zd_close', 'hz_close'])
    
    # 筛选数值型列
    numeric_cols = df_merged.select_dtypes(include=['number']).columns.tolist()
    
    # 筛选合适的自变量（根据问题需求调整）
    excluded = ['zd_close', 'ts_code', 'trade_date', 'Weekday', 'Cash_Flow', 'Index_Value']
    predictors = [col for col in numeric_cols if col not in excluded]
    
    # 检查predictors是否非空
    if not predictors:
        print("警告：没有可用的预测变量。")
        return
    
    # 构建回归模型
    formula = 'zd_close ~ ' + ' + '.join(predictors)
    try:
        results1 = smf.ols(formula, data=df_merged).fit()
        print(results1.summary())
    except Exception as e:
        print(f"构建回归模型时发生错误：{e}")
        return
    
    # 4. 可视化
    # 绘制拟合值与观测值的散点图
    plt.figure(figsize=(10, 6))
    plt.scatter(results1.fittedvalues, df_merged['zd_close'])
    plt.title('Fitted Values vs. Observed Values')
    plt.xlabel('Fitted Values')
    plt.ylabel('Observed Values')
    plt.savefig('Fitted_vs_Observed.png')
    plt.close()
    
    # 计算特征列的相关矩阵（选择部分列进行演示）
    features = df_merged[['Cash_Flow', 'Index_Value', 'hz_close']]
    correlation_matrix = features.corr()
    
    # 绘制相关矩阵的热力图
    plt.figure(figsize=(10, 6))
    sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', cbar=True, fmt='.2f', linewidths=.5)
    plt.title('Correlation Matrix between Features')
    plt.savefig('Correlation_Matrix.png')
    plt.close()
    
    print("数据分析、回归建模和可视化已完成。")

if __name__ == "__main__":
    # 生成Excel文件
    generate_financial_data()
    # 分析Excel文件
    analyze_financial_data()